#include <mambo.hpp>
#include <mpi/communicator.h>
#include <mpi/collectives.h>
#include <cstdlib>
#include <iostream>
#include <fft/fft.h>
using namespace std;

/* driver to test fft libray
 *
 * Need compiler macro: -DMYFFT=0 
 * * -DMYFFT=0 with fftw
 * * -DMYFFT=1 with essl
 * * -DMYFFT=2 with mkl
 */
int main(int argc, char** argv)
{
  mpi::communicator mycomm(argc,argv);

  ///1D FFT dimension
  int fft_dim=12288;
  ///x grid
  int npx=512;
  ///y grid
  int npy=512;
  ///number of ffts to be performed per thread
  int num_fft_do=2;
  ///number of iterations
  int niters=1;
  int ic=0;
  while(ic<argc)
  {
    string a(argv[ic]);
    if(a.find("-s")<a.size()) //get fft dimension
      fft_dim=atoi(argv[++ic]);
    else if(a.find("-npx")<a.size()) //get how many fft to perform
      npx=atoi(argv[++ic]);
    else if(a.find("-npy")<a.size()) //get how many fft to perform
      npy=atoi(argv[++ic]);
    else if(a.find("-d")<a.size()) //number of fftws
      num_fft_do=atoi(argv[++ic]);
    else if(a.find("-i")<a.size()) 
      niters=atoi(argv[++ic]);
    ++ic;
  }

#if defined(FFT_R2C)
  int num_fft=fft_dim/npx*fft_dim/npy;
#else
  int num_fft=fft_dim/npx*fft_dim/npy/2;
#endif
  num_fft_do=std::min(num_fft_do,num_fft);

  double flops_c=(log(static_cast<double>(fft_dim))/log(2.))*5.*fft_dim;
  double flops=flops_c;

  //typedefs for the real and complex
  typedef double real_type;
  typedef complex<double> complex_type;

  ///FFT engine which can do multiple ffts of domany==true
  const bool domany=true;
  int np=omp_get_max_threads();

  mycomm.log() << "BW FFT dim= " << fft_dim 
    << " nx_loc " << fft_dim/npx
    << " ny_loc " << fft_dim/npy
    << " nz_loc " << fft_dim
    << "\nBW target ffts = " <<num_fft
    << " fft done = " << num_fft_do
    << " iterations= " << niters ;

#if defined(FFT_R2C)
  //real-to-complex transforms
  flops*=0.5;
  typedef FFT_DFT<real_type,real_type,domany,MYFFT> fft_eng_type;
  int fft_buffer=fft_dim+2;

  vector<real_type*> target(np);
  vector<fft_eng_type*> fft_doer(np);
#pragma omp parallel
  {
    int ip=omp_get_thread_num();
#if defined(USE_MALLOC)
    target[ip]=static_cast<real_type*>(malloc(fft_buffer*num_fft*sizeof(real_type)));
#else
    target[ip]=new real_type[fft_buffer*num_fft];
#endif
    fft_doer[ip]=new fft_eng_type(fft_dim,target[ip],num_fft_do);
  }
  mycomm.log() << " FLOPS r2c= " << flops<< endl;
#else
  //complex-to-complex transforms
  typedef FFT_DFT<real_type,complex_type,domany,MYFFT> fft_eng_type;
  int fft_buffer=fft_dim;

  vector<complex_type*> target(np);
  vector<fft_eng_type*> fft_doer(np);
#pragma omp parallel
  {
    int ip=omp_get_thread_num();
#if defined(USE_MALLOC)
    target[ip]=static_cast<complex_type*>(malloc(2*fft_buffer*num_fft*sizeof(real_type)));
#else
    target[ip]=new complex_type[fft_buffer*num_fft];
#endif
    fft_doer[ip]=new fft_eng_type(fft_dim,target[ip],num_fft_do);
  }
  mycomm.log() << " FLOPS c2c= " << flops<< endl;
#endif

  mycomm.barrier();
  if(mycomm.rank()==0) cout << "PHASEC" << endl;

  Timer wtimer;
//#pragma omp parallel  firstprivate(num_fft,num_fft_do,fft_buffer)
#pragma omp parallel 
  {
    int ip=omp_get_thread_num();
    int ntot=num_fft*fft_buffer;
    for(int iter=0,offset=0; iter<niters; ++iter)
    {
      fft_doer[ip]->transform_forward(target[ip]+offset,num_fft_do);
      fft_doer[ip]->transform_backward(target[ip]+offset,num_fft_do);
      offset += num_fft_do*fft_buffer;
      if(offset>=ntot) offset -= ntot;
    }
  }
  double dt=wtimer.elapsed();
  mycomm.barrier();
  if(mycomm.rank() == 0) cout << "PHASED" << endl;

  //reduce(mycomm,dt);

  flops *= num_fft_do*2*niters*1e-9;
  //dt /= static_cast<double>(mycomm.size());

  cout << "BW MPI_TASKS = " << mycomm.size()
    <<   " OMP_NUM_THREADS = " << omp_get_max_threads() << endl;
  cout << "BW avgGFLOPS " << flops/dt
    <<" totGFLOPS " << flops << " wtime " << dt << endl;

  //cleanup
  for(int ip=0; ip<np;++ip) delete target[ip];

  return 0;
}

